import os
import numpy as np
import matplotlib.pyplot as plt

import sys
sys.path.insert(0, '/media/tan/_dde_data/projects/ginkgo_vfife/scripts')
from visualization import ResultFileSystem, TimeHistory


def compare(data, keys, figname):
    """ compare results of theroy and openVFIFE

    Args:
        data (dict): {"c": [t, x]}, ..., }
        keys (list): ["c1", "c2", ...]
    """

    figsize = np.array([12, 6]) / 2.54
    fig, ax = plt.subplots(figsize=figsize, tight_layout=True)

    for key in keys:
        print(key, np.mean(data[key][:,1]))
        ax.plot(data[key][:,0], data[key][:,1], lw=1.0, label=key)
    ax.axhline(100, xmin=0.0, xmax=1.0, c="red", lw=1.5, ls="dashed",
               label="theory")
    ax.legend(loc="upper right")
    plt.savefig(figname, dpi=300)
    plt.close(fig)


if __name__ == "__main__":

    # load data from openVFIFE result, differenct damping coeff
    dir = "/home/tan/Desktop/test/vfife_examples/example2/"
    clist = [0.5 * i for i in range(3)]
    keys = ["%.1f" %(i) for i in clist]

    # data container
    damping = {}
    for key in keys:
        fname = os.path.join(dir, "c" + key)
        file = ResultFileSystem(fname)
        th = TimeHistory(file)
        damping[key] = th.extract_link_force(1)

    # save reults
    cwd = os.getcwd()
    figname = os.path.join(cwd, "damping.svg")
    compare(damping, keys, figname)

    # different time step
    # tlist = [0.001, 0.01, 0.1, 0.5]
    tlist = [0.001, 0.01]
    keys = [str(i) for i in tlist]
    times = {}
    for key in keys:
        fname = os.path.join(dir, "t" + key)
        file = ResultFileSystem(fname)
        th = TimeHistory(file)
        times[key] = th.extract_link_force(1)

    # save reults
    cwd = os.getcwd()
    figname = os.path.join(cwd, "timestep.svg")
    compare(times, keys, figname)